import os
import json
import traceback
from typing import Dict, List, Any, Optional, Literal
from datetime import datetime
from dataclasses import dataclass

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_openai import ChatOpenAI

from database.connection import DatabaseClient
from utils.logger import logger
from utils.config import settings
from utils.tools import remove_fk_markup
from .system_prompt import EDIT_SQL_STR, GENERATE_REPORT_STR


@dataclass
class ReportConfig:
    """报告配置"""
    report_type: Literal["html", "markdown"] = "markdown"
    # include_summary: bool = True
    # include_details: bool = True
    # include_recommendations: bool = True


class WorkHoursReportService:
    """工时报告AI服务"""
    
    def __init__(self, api_key: str = None, base_url: str = None):
        """
        初始化报告服务
        
        Args:
            api_key: Deepseek API密钥，如果为None则使用配置文件中的值
            base_url: API基础URL，如果为None则使用配置文件中的值
        """
        # 使用配置文件中的值作为默认值
        self.api_key = api_key or settings.DEEPSEEK_API_KEY
        # 直接使用配置文件中的基础URL
        self.base_url = base_url or settings.DEEPSEEK_BASE_URL
        self.db_connection = DatabaseClient()
        
        # 验证必要的配置
        if not self.api_key:
            raise ValueError("DEEPSEEK_API_KEY不能为空，请在环境变量中设置或通过参数传入")
        
        # 初始化LLM
        self.llm = ChatOpenAI(
            model=settings.DEEPSEEK_MODEL,
            openai_api_key=self.api_key,
            openai_api_base=self.base_url,
            temperature=settings.DEEPSEEK_TEMPERATURE,
            max_tokens=settings.DEEPSEEK_MAX_TOKENS
        )
        
        # 创建SQL编辑链
        self.sql_edit_chain = self._create_sql_edit_chain()
        
        # 创建报告生成链
        self.report_chain = self._create_report_chain()
        
        logger.info("🤖 WorkHoursReportService initialized successfully")
    
    def _create_sql_edit_chain(self) -> LLMChain:
        """创建SQL编辑链"""
        prompt = PromptTemplate(
            input_variables=["user_query"],
            template=EDIT_SQL_STR + "\n\n用户问题: {user_query}"
        )
        
        return LLMChain(
            llm=self.llm,
            prompt=prompt,
            verbose=True
        )
    
    def _create_report_chain(self) -> LLMChain:
        """创建报告生成链"""
        prompt = PromptTemplate(
            input_variables=["data", "report_type"],
            template=GENERATE_REPORT_STR
        )
        
        return LLMChain(
            llm=self.llm,
            prompt=prompt,
            verbose=True
        )
    
    def generate_sql_from_query(self, user_query: str) -> Dict[str, str]:
        """
        根据用户问题生成SQL查询
        
        Args:
            user_query: 用户的自然语言查询
            
        Returns:
            包含SQL语句和状态信息的字典
        """
        try:
            logger.info(f"🔍 Processing user query: {user_query}")
            
            # 使用LLM生成SQL
            response = self.sql_edit_chain.invoke({"user_query": user_query})
            
            # 处理invoke返回的结果
            if isinstance(response, dict) and "text" in response:
                response_text = response["text"]
            elif isinstance(response, dict) and "output" in response:
                response_text = response["output"]
            else:
                response_text = str(response)
            
            # 解析JSON响应
            try:
                logger.info(f"🔍 Response: {remove_fk_markup(response_text)}")
                result = json.loads(remove_fk_markup(response_text))
                logger.info(f"📝 Generated SQL: {result}")
                return result
            except json.JSONDecodeError:
                logger.error(f"❌ Failed to parse LLM response as JSON: {remove_fk_markup(response_text)}")
                return {"sql": "ERROR", "message": "LLM响应格式错误"}
                
        except Exception as e:
            logger.error(f"❌ Error generating SQL: {e}")
            return {"sql": "ERROR", "message": f"生成SQL时发生错误: {str(e)}"}
    
    def execute_query_with_ai(self, user_query: str) -> Dict[str, Any]:
        """
        使用AI执行查询
        
        Args:
            user_query: 用户的自然语言查询
            
        Returns:
            查询结果
        """
        try:
            # 生成SQL
            sql_result = self.generate_sql_from_query(user_query)
            
            if sql_result["sql"] == "ERROR":
                return {
                    "success": False,
                    "message": sql_result["message"],
                    "data": None
                }
            
            # 执行SQL查询
            sql_query = sql_result["sql"]
            
            # 检查是否为SELECT查询
            if not sql_query.strip().upper().startswith('SELECT'):
                return {
                    "success": False,
                    "message": "错误：只允许SELECT查询语句",
                    "data": None
                }
            
            # 执行查询
            try:
                results = self.db_connection.execute_query(sql_query)
                return {
                    "success": True,
                    "message": "查询成功",
                    "data": results,
                    "sql": sql_query
                }
            except Exception as e:
                error_msg = f"SQL查询执行失败: {str(e)}"
                logger.error(error_msg)
                return {
                    "success": False,
                    "message": error_msg,
                    "data": None
                }
                
        except Exception as e:
            logger.error(f"❌ Error in execute_query_with_ai: {e}")
            return {
                "success": False,
                "message": f"查询执行失败: {str(e)}",
                "data": None
            }
    
    def generate_report(
        self,
        data: List[Dict[str, Any]],
        config: ReportConfig = ReportConfig()
    ) -> Dict[str, Any]:
        """
        生成工时报告
        
        Args:
            data: 工时数据
            config: 报告配置
            
        Returns:
            生成的报告
        """
        try:
            logger.info(f"📊 Generating {config.report_type} report for {len(data)} records")
            
            # 准备数据
            data_str = json.dumps(data, ensure_ascii=False, indent=2, default=str)
            
            # 生成报告
            response = self.report_chain.invoke({
                "data": data_str,
                "report_type": config.report_type
            })
            
            # 处理invoke返回的结果
            if isinstance(response, dict) and "text" in response:
                report_content = response["text"]
            elif isinstance(response, dict) and "output" in response:
                report_content = response["output"]
            else:
                report_content = str(response)
            
            # 添加统计信息
            total_hours = sum(float(record.get('WorkHour', 0)) for record in data)
            total_records = len(data)
            report_content = remove_fk_markup(report_content, config.report_type)
            
            return {
                "success": True,
                "report_type": config.report_type,
                "content": report_content,
                "statistics": {
                    "total_hours": total_hours,
                    "total_records": total_records,
                    "generated_at": datetime.now().isoformat()
                }
            }
            
        except Exception as e:
            logger.error(f"❌ Error generating report: {e} - {traceback.format_exc()}")
            return {
                "success": False,
                "message": f"生成报告失败: {str(e)}",
                "content": None
            }
    
    def _format_html_report(self, content: str, total_hours: float, total_records: int, data_str: dict) -> str:
        """格式化HTML报告"""
        # 解析data_str为JSON数据
        try:
            data = json.loads(data_str)
        except json.JSONDecodeError:
            data = []
        
        # 生成表格HTML
        table_html = self._generate_data_table(data)
        
        html_template = content + f"""
            <div class="header">
                <h1>🐂 牛马工时报告</h1>
                <p>生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
            </div>
            
            <div class="statistics">
                <h2>📊 统计信息</h2>
                <p><strong>总工时:</strong> {total_hours:.2f} 小时</p>
                <p><strong>记录数:</strong> {total_records} 条</p>
            </div>
            <div class="data-table">
                <h2>📋 元数据详情</h2>
                {table_html}
            </div>
        """
        return html_template
    
    def _generate_data_table(self, data: List[Dict[str, Any]]) -> str:
        """生成数据表格HTML"""
        if not data:
            return "<p>暂无数据</p>"
        
        # 获取所有字段名作为表头
        all_keys = set()
        for record in data:
            all_keys.update(record.keys())
        
        headers = sorted(list(all_keys))
        
        # 生成表格HTML
        table_html = """
        <table border="1" style="border-collapse: collapse; width: 100%; margin-top: 10px;">
            <thead>
                <tr style="background-color: #f0f0f0;">
        """
        
        # 添加表头
        for header in headers:
            table_html += f"<th style='padding: 8px; text-align: left;'>{header}</th>"
        
        table_html += """
                </tr>
            </thead>
            <tbody>
        """
        
        # 添加数据行
        for record in data:
            table_html += "<tr>"
            for header in headers:
                value = record.get(header, "")
                table_html += f"<td style='padding: 8px; border: 1px solid #ddd;'>{value}</td>"
            table_html += "</tr>"
        
        table_html += """
            </tbody>
        </table>
        """
        
        return table_html
    
    def _format_markdown_report(self, content: str, total_hours: float, total_records: int, data_str: dict) -> str:
        """格式化Markdown报告"""
        markdown_template = content + f"""# 🐂 牛马工时源数据资料

**生成时间:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

## 📊 统计信息

- **总工时:** {total_hours:.2f} 小时
- **记录数:** {total_records} 条

## 📝 原始数据
```json
{data_str}
```

---
*本报告由AI智能生成*
"""
        return markdown_template
    
    def generate_comprehensive_report(
        self,
        user_query: str,
        config: ReportConfig = ReportConfig()
    ) -> Dict[str, Any]:
        """
        生成综合报告（查询+分析）
        
        Args:
            user_query: 用户查询
            config: 报告配置
            
        Returns:
            综合报告结果
        """
        try:
            logger.info(f"🔄 Generating comprehensive report for query: {user_query}")
            
            # 1. 执行查询
            query_result = self.execute_query_with_ai(user_query)
            
            if not query_result["success"]:
                return query_result
            
            # 2. 生成报告
            report_result = self.generate_report(query_result["data"], config)
            
            if not report_result["success"]:
                return report_result
            
            # 3. 合并结果
            return {
                "success": True,
                "query": user_query,
                "sql": query_result.get("sql"),
                "data": query_result["data"],
                "report": report_result["content"],
                "report_type": config.report_type,
                "statistics": report_result["statistics"]
            }
            
        except Exception as e:
            logger.error(f"❌ Error generating comprehensive report: {e}")
            return {
                "success": False,
                "message": f"生成综合报告失败: {str(e)}"
            }


# 创建全局服务实例（可选）
_report_service = None

def get_report_service(api_key: str = None, base_url: str = None) -> WorkHoursReportService:
    """获取报告服务实例"""
    global _report_service
    
    if _report_service is None:
        # 如果没有传入参数，使用配置文件中的值
        if not api_key and not settings.DEEPSEEK_API_KEY:
            raise ValueError("请设置DEEPSEEK_API_KEY环境变量或通过参数传入api_key")
        
        _report_service = WorkHoursReportService(api_key, base_url)
    
    return _report_service 